"""
Adapted from salesforce@LAVIS. Below is the original copyright:
 Copyright (c) 2022, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""

import logging
import os

import numpy as np
import torch
import torch.nn as nn
from timechat.common.dist_utils import download_cached_file, is_dist_avail_and_initialized
from timechat.common.utils import get_abs_path, is_url
from omegaconf import OmegaConf


class BaseModel(nn.Module):
    """Base class for models."""

    def __init__(self):
        super().__init__()

    @property
    def device(self):
        return list(self.parameters())[0].device  #返回模型第一个参数所在的设备

    def load_checkpoint(self, url_or_filename):
        """
        Load from a finetuned checkpoint.

        This should expect no mismatch in the model keys and the checkpoint keys.
        """

        if is_url(url_or_filename):
            cached_file = download_cached_file(
                url_or_filename, check_hash=False, progress=True
            ) #检查输入是否为URL，下载URL对应的文件并缓存
            checkpoint = torch.load(cached_file, map_location="cpu") #加载缓存文件中的检查点到cpu
        elif os.path.isfile(url_or_filename):
            checkpoint = torch.load(url_or_filename, map_location="cpu") #检查输入是否为文件路径，加载文件路径中的检查点到CPU
        else:
            raise RuntimeError("checkpoint url or path is invalid")

        if "model" in checkpoint.keys():
            state_dict = checkpoint["model"] #检查检查点是否包含“model”键，如有，获取对应的值作为状态字典
        else:
            state_dict = checkpoint #如果不包含model键，直接将检查点作为状态字典

        msg = self.load_state_dict(state_dict, strict=False) #加载状态字典到模型中，strict=False允许非严格匹配

        logging.info("Missing keys {}".format(msg.missing_keys)) #记录缺失的键信息
        logging.info("load checkpoint from %s" % url_or_filename) #记录加载检查点的信息

        return msg #返回加载状态信息

    @classmethod
    def from_pretrained(cls, model_type):
        """
        Build a pretrained model from default configuration file, specified by model_type.

        Args:
            - model_type (str): model type, specifying architecture and checkpoints.

        Returns:
            - model (nn.Module): pretrained or finetuned model, depending on the configuration.
        """
        model_cfg = OmegaConf.load(cls.default_config_path(model_type)).model #加载默认配置文件中的模型配置
        model = cls.from_config(model_cfg) #根据配置创建模型实例

        return model

    @classmethod
    def default_config_path(cls, model_type): #获取默认配置文件路径的类方法
        assert (
            model_type in cls.PRETRAINED_MODEL_CONFIG_DICT
        ), "Unknown model type {}".format(model_type) #断言模型类型在预定义的配置字典中
        return get_abs_path(cls.PRETRAINED_MODEL_CONFIG_DICT[model_type]) #返回模型类型对应的配置文件绝对路径

    #从配置加载检查点的方法
    def load_checkpoint_from_config(self, cfg, **kwargs):
        """
        Load checkpoint as specified in the config file.

        If load_finetuned is True, load the finetuned model; otherwise, load the pretrained model.
        When loading the pretrained model, each task-specific architecture may define their
        own load_from_pretrained() method.
        """
        load_finetuned = cfg.get("load_finetuned", True) #获取是否加载微调模型的配置，默认为true
        if load_finetuned:
            finetune_path = cfg.get("finetuned", None) #获取微调模型路径
            assert (
                finetune_path is not None
            ), "Found load_finetuned is True, but finetune_path is None." #断言微调路径不为空
            self.load_checkpoint(url_or_filename=finetune_path) #加载微调检查点
        else:
            # load pre-trained weights
            pretrain_path = cfg.get("pretrained", None) #获取预训练模型路径
            assert "Found load_finetuned is False, but pretrain_path is None." #断言预训练路径不为空
            self.load_from_pretrained(url_or_filename=pretrain_path, **kwargs) #加载预训练模型

    def before_evaluation(self, **kwargs):
        pass

    #定义显示模型参数数量的方法
    def show_n_params(self, return_str=True):
        tot = 0 #初始化总参数数量为0
        for p in self.parameters(): #遍历模型所有参数
            w = 1 #初始化参数大小为1
            for x in p.shape: #遍历参数形状每个维度
                w *= x #计算参数的总元素数量
            tot += w #累加到总参数数量
        if return_str: #是否以字符串形式返回
            if tot >= 1e6: #总参数是否大于等于一百万
                return "{:.1f}M".format(tot / 1e6) #以百万为单位返回参数数量字符串
            else:
                return "{:.1f}K".format(tot / 1e3) #以千为单位返回
        else:
            return tot

#基础编码器
class BaseEncoder(nn.Module):
    """
    Base class for primitive encoders, such as ViT, TimeSformer, etc.
    """

    def __init__(self):
        super().__init__()

    def forward_features(self, samples, **kwargs):
        raise NotImplementedError #抛出未实现错误，提示子类需要实现该方法

    @property
    def device(self):
        return list(self.parameters())[0].device #返回模型第一个参数所在的设备

#混合类，用于处理队列操作
class SharedQueueMixin:
    @torch.no_grad() #表示该方法在执行时不计算梯度
    def _dequeue_and_enqueue(self, image_feat, text_feat, idxs=None): #出队和入队
        # gather keys before updating queue
        image_feats = concat_all_gather(image_feat) #聚合所有进程中的图像特征
        text_feats = concat_all_gather(text_feat) #聚合所有进程中的文本特征

        batch_size = image_feats.shape[0] #获取批量大小

        ptr = int(self.queue_ptr) #获取当前指针所在位置
        assert self.queue_size % batch_size == 0  # for simplicity 断言队列大小能被批量大小整除

        # replace the keys at ptr (dequeue and enqueue)
        self.image_queue[:, ptr : ptr + batch_size] = image_feats.T #将图像特征入队
        self.text_queue[:, ptr : ptr + batch_size] = text_feats.T #将文本特征入队

        if idxs is not None: #如果索引不为空
            idxs = concat_all_gather(idxs) #聚合进程中的所有索引
            self.idx_queue[:, ptr : ptr + batch_size] = idxs.T #将索引入队

        ptr = (ptr + batch_size) % self.queue_size  # move pointer 更新指针位置
        self.queue_ptr[0] = ptr #更新队列指针

#混合类，用于动量蒸馏
class MomentumDistilationMixin:
    @torch.no_grad()
    def copy_params(self):
        for model_pair in self.model_pairs: #遍历模型对
            for param, param_m in zip(
                model_pair[0].parameters(), model_pair[1].parameters()
            ): #遍历模型对中的参数
                param_m.data.copy_(param.data)  # initialize 将参数数据复制到动量模型
                param_m.requires_grad = False  # not update by gradient 禁用梯度计算

    @torch.no_grad()
    def _momentum_update(self): #动量更新
        for model_pair in self.model_pairs:
            for param, param_m in zip(
                model_pair[0].parameters(), model_pair[1].parameters()
            ):
                param_m.data = param_m.data * self.momentum + param.data * (
                    1.0 - self.momentum
                ) #动量更新公式

#张量聚合
class GatherLayer(torch.autograd.Function):
    """
    Gather tensors from all workers with support for backward propagation:
    This implementation does not cut the gradients as torch.distributed.all_gather does.
    """

    @staticmethod
    def forward(ctx, x):
        output = [
            torch.zeros_like(x) for _ in range(torch.distributed.get_world_size())
        ] #初始化输出列表
        torch.distributed.all_gather(output, x) #聚合所有进程中的张量
        return tuple(output) #返回聚合后的张量元组

    @staticmethod
    def backward(ctx, *grads):
        all_gradients = torch.stack(grads) #将梯度堆叠成一个张量
        torch.distributed.all_reduce(all_gradients) #对梯度进行全归约操作
        return all_gradients[torch.distributed.get_rank()] #返回当前进程的梯度

#用于带有梯度的全聚集操作
def all_gather_with_grad(tensors):
    """
    Performs all_gather operation on the provided tensors.
    Graph remains connected for backward grad computation.
    """
    # Queue the gathered tensors
    world_size = torch.distributed.get_world_size() #获取进程数
    # There is no need for reduction in the single-proc case
    if world_size == 1:
        return tensors #单进程时不聚集

    # tensor_all = GatherLayer.apply(tensors)
    tensor_all = GatherLayer.apply(tensors) #应用GatherLayer进行聚集

    return torch.cat(tensor_all, dim=0) #按维度0拼接张量

#用于全聚集操作
@torch.no_grad()
def concat_all_gather(tensor):
    """
    Performs all_gather operation on the provided tensors.
    *** Warning ***: torch.distributed.all_gather has no gradient.
    """
    # if use distributed training
    if not is_dist_avail_and_initialized():
        return tensor #非分布式训练时直接返回张量

    tensors_gather = [
        torch.ones_like(tensor) for _ in range(torch.distributed.get_world_size())
    ] #初始化聚集张量列表
    torch.distributed.all_gather(tensors_gather, tensor, async_op=False) #聚合所有进程中的张量

    output = torch.cat(tensors_gather, dim=0) #按维度0进行拼接张量
    return output

#用于在指定维度上复制张量
def tile(x, dim, n_tile):
    init_dim = x.size(dim) #获取张量在指定维度上的大小
    repeat_idx = [1] * x.dim() #初始化重复索引列表
    repeat_idx[dim] = n_tile #设置指定维度的重复次数
    x = x.repeat(*(repeat_idx)) #复制张量
    order_index = torch.LongTensor(
        np.concatenate([init_dim * np.arange(n_tile) + i for i in range(init_dim)])
    ) #创建索引张量
    return torch.index_select(x, dim, order_index.to(x.device)) #根据索引选择元素并返回
